package com.sts.sink.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.poi.excel.BigExcelWriter;
import cn.hutool.poi.excel.ExcelUtil;
import com.sts.sink.StsSink;
import com.sts.sink.model.ExcelSinkConfig;
import com.sts.source.StsSource;
import com.sts.task.StsTask;
import com.sts.transform.StsTransform;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
import java.util.stream.Gatherers;
import java.util.stream.Stream;

/**
 * @author lsl
 * @since 2025年06月28日 13:48
 */
@Slf4j
public class ExcelSinkImpl implements StsSink {

    private final StsSource stsSource;
    private final ExcelSinkConfig excelSinkConfig;
    private final StsTask stsTask;

    public ExcelSinkImpl(StsSource stsSource, ExcelSinkConfig excelSinkConfig, StsTask stsTask) {
        this.stsSource = stsSource;
        this.excelSinkConfig = excelSinkConfig;
        this.stsTask = stsTask;
    }

    @Override
    public void insert() {
        try {
            // 如果存在就先删除, 否则会导致程序oom
            Path path = Paths.get(excelSinkConfig.getFilePath());
            if (Files.exists(path)) {
                Files.delete(path);
            }
            try (BigExcelWriter writer = ExcelUtil.getBigWriter(excelSinkConfig.getFilePath(), excelSinkConfig.getSheetName())) {
                int limit = 1048576;
                if (excelSinkConfig.isHasHeader()) {
                    // 导出头
                    List<String> header = stsSource.getHeader();
                    writer.writeHeadRow(header);
                    limit = 1048575;
                }

                Stream<List<String>> dataStream = stsSource.getDataStream();

                if (CollectionUtil.isNotEmpty(stsTask.getStsTransformList())) {
                    for (StsTransform<?> stsTransform : stsTask.getStsTransformList()) {
                        dataStream = stsTransform.transform(dataStream);
                    }
                }

                dataStream = dataStream.limit(limit);

                dataStream.takeWhile(_ -> stsTask.isFlag())
                        .gather(Gatherers.windowFixed(excelSinkConfig.getBatchSize()))
                        .forEach(dataList -> {
                            writer.write(dataList);
                            stsTask.updateTotalSuccessCount(stsTask, dataList.size());
                        });
            }
        } catch (IOException e) {
            throw new RuntimeException(e.getMessage(), e);
        } finally {
            try {
                stsSource.close();
            } catch (Exception e) {
                log.error("关闭连接失败: {}", e.getMessage(), e);
            }
        }
    }

    @Override
    public void upsert() {
        insert();
    }
}
